package com.socket2;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.lang.Assert;
import cn.hutool.core.lang.Console;
import cn.hutool.core.util.ArrayUtil;
import cn.hutool.core.util.StrUtil;

import java.io.*;
import java.net.InetAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.util.concurrent.CountDownLatch;

public class Socket5Proxy {


    private static final byte TYPE_IPV4 = 0x01;
    private static final byte TYPE_IPV6 = 0X02;
    private static final byte TYPE_HOST = 0X03;


    public static void main(String[] args) throws IOException {
        ServerSocket serverSocket = new ServerSocket(11011);
        while (true) {
            try {
                Console.log("lister in "+serverSocket.getLocalPort());
                Socket socket = serverSocket.accept();
                Console.log("connect success..");
                InputStream inputStream = socket.getInputStream();
                DataInputStream sourceIn = new DataInputStream(inputStream);

                OutputStream outputStream = socket.getOutputStream();
                DataOutputStream sourceOut = new DataOutputStream(outputStream);
                int i;
                while ((i = sourceIn.read()) != -1) {
                    // 代表 SOCKS 的版本，SOCKS5 默认为0x05，其固定长度为1个字节；
                    Console.log("VER:{}",i);
                    // 表示第三个字段METHODS的长度，它的长度也是1个字节；
                    int authSize = sourceIn.read();
                    Console.log("AUTH SIZE:{}", authSize);
                    // 表示客户端支持的验证方式，可以有多种，他的长度是1-255个字节。
                    byte[] bytes = new byte[authSize];
                    sourceIn.read(bytes);

                    // 为了简化demo我们服务端目前只处理最简单 不认证 代理方式
                    if (ArrayUtil.contains(bytes, (byte)0)) {
                        System.out.println("support no auth");
                    }
                    // 上面是已经接受了客户端全部的 建立连接请求包. 接下来要服务端返回一个响应了
                    // 0x05代表是SOCKS5版本协议, 0x00代表不需要认证
                    sourceOut.write(new byte[]{0x05, 0x00});
                    sourceOut.flush();

                    // 客户端也接受成后, 开始发送他想访问的目标服务器信息
                    // 这里一次性读取4个字节, 因为这里有四个字段,且每个的长度固定为1个字节, 我就直接读取了
                    bytes = new byte[4];
                    sourceIn.read(bytes);
                    // 代表 SOCKS 协议的版本，SOCKS 默认为0x05，其值长度为1个字节；
                    byte VER = bytes[0];
                    // 代表客户端请求的类型，值长度也是1个字节，有三种类型；
                    byte CMD = bytes[1];
                    // 保留字，值长度为1个字节；
                    byte RSV = bytes[2];
                    // 代表请求的远程服务器地址类型，值长度1个字节，有三种类型；
                    byte ATYP = bytes[3];

                    // 本demo只测试SOCKS5版本, CONNECT类型, IP访问的请求
                    Assert.isTrue(VER == 0x05);
                    Assert.isTrue(CMD == 0x01);
                    Assert.isTrue(ATYP == 0x01);

                    // 根据地址类型字段, 构造出目标地址的host, 可能为x.x.x.x(ip4) 也可能为 xxx.com(domain)
                    String host = getHost(ATYP, sourceIn);
                    Console.log("read host:{}", host);
                    // 代表远程服务器的端口，要访问哪个端口的意思，值长度2个字节。
                    bytes = new byte[2];
                    sourceIn.read(bytes);
                    // 转成10进制表示端口号
                    short port = ByteBuffer.wrap(bytes).getShort();

                    // 与目标服务器建立连接
                    Socket proxy = new Socket(host, port);
                    // 建立成功后,通知客户端, 按照协议规定的通知
                    bytes = new byte[]{0x05,0x00,0x00,0x01};
                    sourceOut.write(bytes);
                    // 这里的地址和端口好像随便填, 不影响
                    sourceOut.write(socket.getLocalAddress().getAddress());
                    Short localPort = (short) ((socket.getLocalPort()) & 0xFFFF);
                    sourceOut.writeShort(localPort);
                    sourceOut.flush();

                    // 这里比较难理解
                    // 同时打开客户端A 到代理服务器P 的输入输出流
                    // 也打开代理服务器P 到目标服务器B 的输入输出流
                    // 即 如果客户端A有内容写入的时候, 同时将写入的数据传递给B
                    // A -> sourceIn -> P -> proxyOut -> B
                    // 同理, 目标服务器B有响应返回的时候, 同时将数据返回给A
                    // B -> proxyIn -> P -> sourceOut -> A
                    OutputStream proxyOut = proxy.getOutputStream();
                    InputStream proxyIn = proxy.getInputStream();
                    System.out.println("server proxy port:"+proxy.getLocalPort());
                    CountDownLatch countDownLatch = new CountDownLatch(1);
                    transfer(sourceIn, proxyOut, countDownLatch);
                    transfer(proxyIn, sourceOut, countDownLatch);
                    // 如果有一方断开, 那就可以终止这次代理了. 因为这代表着有一方已经4次挥手啦
                    try {
                        countDownLatch.await();
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }
                    proxy.close();
                    Console.log("proxy end");
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    private static String getHost(byte type, InputStream inputStream) throws IOException {
        String host = null;
        byte[] tmp = null;
        switch (type) {
            case TYPE_IPV4:
                tmp = new byte[4];
                inputStream.read(tmp);
                host = InetAddress.getByAddress(tmp).getHostAddress();
                break;
            case TYPE_IPV6:
                tmp = new byte[16];
                inputStream.read(tmp);
                host = InetAddress.getByAddress(tmp).getHostAddress();
                break;
            case TYPE_HOST:
                int count = inputStream.read();
                tmp = new byte[count];
                inputStream.read(tmp);
                host = new String(tmp);
            default:
                break;
        }
        return host;
    }

    private static void transfer(InputStream in, OutputStream out, CountDownLatch latch) {
        new Thread(() -> {
            byte[] bytes = new byte[1024];
            int count = 0;
            try {
                while (0 < (count = in.read(bytes))) {
                    out.write(bytes, 0, count);
                    out.flush();
                }
            } catch (IOException e) {
                Console.error(e, "转换出现错误");
                e.printStackTrace();
            }
            if (latch != null) {
                latch.countDown();
            }

        }).start();

    }
}
